# coding: UTF-8
import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn import metrics
from transformers import optimization
from utils.utils import get_time_dif


def train(config, model, train_iter, val_iter, test_iter, loss_function=F.cross_entropy):
    start_time = time.time()
    model.train()
    param_optimizer = list(model.named_parameters())
    no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]
    learning_rate = config.learning_rate
    optimizer = optimization.AdamW(optimizer_grouped_parameters, lr=learning_rate)

    total_batch = 0  # 记录进行到多少batch
    val_best_loss = float('inf')
    last_improve = 0  # 记录上次验证集loss下降的batch数
    flag = False  # 记录是否很久没有效果提升
    for epoch in range(config.num_epochs):
        print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))
        for i, (texts, labels) in enumerate(train_iter):
            #print(labels)
            outputs = model(texts)
            model.zero_grad()
            loss = loss_function(outputs, labels)
            loss.backward()
            optimizer.step()
            if total_batch % 10 == 0:
                # 每多少轮输出在训练集和验证集上的效果
                true = labels.data.cpu()
                predic = torch.max(outputs.data, 1)[1].cpu()
                train_acc = metrics.accuracy_score(true, predic)
                val_acc, val_loss = evaluate(config, model, val_iter, loss_function)
                if val_loss < val_best_loss:
                    val_best_loss = val_loss
                    torch.save(model.state_dict(), config.save_path)
                    improve = '*'
                    last_improve = total_batch
                else:
                    improve = ''
                time_dif = get_time_dif(start_time)
                msg = 'Iter: {0:>6},  Train Loss: {1:>5.2},  Train Acc: {2:>6.2%},  Val Loss: {3:>5.2},  Val Acc: {4:>6.2%},  Time: {5} {6}'
                print(msg.format(total_batch, loss.item(), train_acc, val_loss, val_acc, time_dif, improve))
                model.train()
            '''
            # 学习率手动衰减， AdamW 本身就是实现了学习率衰减的学习器
            if total_batch % 50 == 0:
                learning_rate /= 10
                optimizer = optimization.AdamW(optimizer_grouped_parameters, lr=learning_rate)
            '''
            total_batch += 1
            if total_batch - last_improve > config.require_improvement:
                # 验证集loss超过1000batch没下降，结束训练
                print("No optimization for a long time, auto-stopping...")
                flag = True
                break
        if flag:
            break
    test(config, model, test_iter, loss_function)

def test(config, model, test_iter, loss_function=F.cross_entropy):
    model.load_state_dict(torch.load(config.save_path, map_location=config.device))
    model.eval()
    start_time = time.time()
    test_acc, test_loss, test_report, test_confusion = evaluate(config, model, test_iter, loss_function, test=True)
    msg = 'Test Loss: {0:>5.2},  Test Acc: {1:>6.2%}'
    print(msg.format(test_loss, test_acc))
    print("Precision, Recall and F1-Score...")
    print(test_report)
    print("Confusion Matrix...")
    print(test_confusion)
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)

def evaluate(config, model, data_iter, loss_function, test=False):
    model.eval()
    loss_total = 0
    predict_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    softmax = nn.Softmax(dim=0)
    with torch.no_grad():
        for texts, labels in data_iter:
            outputs = model(texts)

            loss = loss_function(outputs, labels)
            loss_total += loss
            labels = labels.data.cpu().numpy()
            
            predict_batch = []
            for i, output in enumerate(outputs):
                output = softmax(output)
                if output[0] - 0.2 > output[1]:
                    l = 0
                else:
                    l = 1
                predict_batch.append(l)
                if labels[i] == 1 and l == 0:
                    print(output)
            labels_all = np.append(labels_all, labels)
            predict_all = np.append(predict_all, predict_batch)
            '''
            for i, label in enumerate(labels):
                if label != predic[i]:
                    print('ground truth:', label)
                    print('predic:', predic[i])
                    print('source:', texts[i])
                    print('- - -')
                    #input()
            '''
            

    acc = metrics.accuracy_score(labels_all, predict_all)
    if test:
        report = metrics.classification_report(labels_all, predict_all, target_names=config.class_list, digits=4)
        confusion = metrics.confusion_matrix(labels_all, predict_all)
        return acc, loss_total / len(data_iter), report, confusion
    return acc, loss_total / len(data_iter)

if __name__ == "__main__":
    print('a')
    a = torch.FloatTensor([[-0.3867,  0.0470, -0.2407,  0.4931],
        [-0.3756,  0.5491, -0.0152,  0.5762],
        [-0.4836,  0.4270,  0.5146,  0.2515],
        [-0.3064,  0.2886, -0.1370,  0.4034],
        [-0.5047,  0.0930,  0.0552,  0.4332]])
    b = torch.LongTensor([0, 1, 3, 3, 0])
    F_loss = FocalLoss(4)
    loss = F_loss(a,b)
